Skip to content

Propagate disable_torch_autograd to thunderfx's _splitter #2534

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Propagate disable_torch_autograd to thunderfx's _splitter #2534

wants to merge 2 commits into from
+20 −9

Conversation

crcrpar
Copy link
Collaborator

@crcrpar crcrpar commented Sep 25, 2025 •
edited
Loading

What does this PR do?

Fixes somewhat unexpected graph splits in thunderfx path with programs that use torch.library.custom_op even when it's manually registered through _register_custom_op.
For example, when there's a torch.library.custom_op that's registered to Thunder without backward definition, that custom_op could result in a graph split. This is because splitter only sees input proxies requires_grad --

lightning-thunder/thunder/dynamo/utils.py

Lines 312 to 322 in 280c57e

function_to_run = value_and_grad(thunder_symbol) if requires_grad else thunder_symbol
# We need to be under trace context to generate proxies.
with thunder.core.trace.tracectx(TraceCtx()):
try:
function_to_run(*proxy_args, **proxy_kwargs)
except Exception as e:
return False, SplitReason(
SplitReasonType.EXCEPTION_META_THUNDER_OP,
f"Failed while running meta for node with name: {node.name} and target: {node.target}, see exception field",
exception=str(e),
)
.

close #2606

This was referenced Sep 25, 2025
Draft
Open
@crcrpar
Propagate disable_torch_autograd to thunderfx's _splitter
ec045fe 
For example, when there's a `torch.library.custom_op` that's registered
to Thunder without backward definition, that custom_op could result in a
graph split. This is because splitter only sees input proxies
`requires_grad` -- https://github.com/Lightning-AI/lightning-thunder/blob/280c57ec289d09eacad0a29cf97a332e02ecebaa/thunder/dynamo/utils.py#L312-L322

Signed-off-by: Masaki Kozuki <[email protected]>
@crcrpar crcrpar force-pushed the respect-disable-torch-autograd-in-thunderfx-splitter branch from 0589d53 to ec045fe Compare October 7, 2025 05:49
@crcrpar crcrpar requested a review from KaelanDt as a code owner October 7, 2025 05:49
shino16
shino16 reviewed Oct 7, 2025
View reviewed changes
thunder/dynamo/utils.py Outdated
)

function_to_run = value_and_grad(thunder_symbol) if requires_grad else thunder_symbol
function_to_run = thunder_symbol
Copy link
Collaborator

@shino16 shino16 Oct 7, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this line redundant?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah yes, great catch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@shino16 shino16 shino16 approved these changes

@mruberry mruberry Awaiting requested review from mruberry mruberry is a code owner

@lantiga lantiga Awaiting requested review from lantiga lantiga is a code owner

@t-vi t-vi Awaiting requested review from t-vi t-vi is a code owner

@KaelanDt KaelanDt Awaiting requested review from KaelanDt KaelanDt is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
2 participants
@crcrpar @shino16